#include <string.h>
#include "sys/queue.h"
#include "esp_flash.h"
#include "esp_flash_partitions.h"
#include "esp_ota_ops.h"
#include "esp_app_format.h"
#include "spi_flash_mmap.h"
#include "esp_log.h"
#include "ota.h"

#define dbus_change_to_process(...)
#define dbus_only_change(...)

static const char *TAG = "OTA/Stream";
#define LOG_TAG "OTA/Stream"

#define LOG_PROGRESS(...)                 \
    do                                    \
    {                                     \
        printf("\r\e[32m[" LOG_TAG "] "); \
        printf(__VA_ARGS__);              \
        printf("\e[0m");                  \
    } while (0)

uint8_t http_stream_header_cb(void *handle, char *key, char *value)
{
    if (!handle)
        handle = ota_default_handle;
    ota_t ota = handle;
    if (!ota)
    {
        dbus_change_to_process("ota", "Missing parameter");
        return 2;
    }
    char *p = strnstr(key, "Content-Length", 15);
    if (!p)
        p = strnstr(key, "content-length", 15);
    if (p)
        ota->length = atol(value);
    p = strnstr(key, "Content-Range", 14);
    if (!p)
        p = strnstr(key, "content-range", 14);
    if (p)
    {
        uint32_t offset = 0;
        sscanf(value, "bytes %lu-", &offset);
        if (offset && offset < ota->length)
            ota->offset = offset;
        else
            ota->offset = 0;
    }
    else
        ota->offset = 0;
    return 0;
}

static uint8_t http_check_header(ota_t ota, uint8_t **s_data, uint32_t *s_len)
{
    uint8_t *data = *s_data;
    uint32_t len = *s_len;
    uint32_t offset = 0;
    char *len_str = strstr((char *)data, "Content-Length:");
    if (!len_str)
    {
        len_str = strstr((char *)data, "content-length:");
        if (!len_str)
        {
            sprintf(ota->err_str, "Missing Length");
            ESP_LOGE(TAG, "%s", ota->err_str);
            dbus_change_to_process("ota", "Missing Length");
            // elog_hexdump("", 16, data, len);
            return 1;
        }
    }
    ota->length = atol(len_str + 15);
    char *range_str = strstr((char *)data, "Content-Range:");
    if (!range_str)
        range_str = strstr((char *)data, "content-range:");
    if (range_str)
    {
        sscanf(range_str + 15, "bytes %lu-", &offset);
        if (offset && offset < ota->length)
            ota->offset = offset;
        else
            ota->offset = 0;
    }
    else
        ota->offset = 0;
    uint8_t *p = (uint8_t *)strstr((char *)data, "\n\n") + 3;
    if (p == 3)
    {
        p = (uint8_t *)strstr((char *)data, "\r\n\r\n") + 4;
        if (p == 4)
        {
            sprintf(ota->err_str, "Missing header");
            ESP_LOGE(TAG, "%s", ota->err_str);
            // dbus_change_to_process("ota", "Missing header");
            // elog_hexdump("", 16, data, len);
            return 2;
        }
    }
    len -= (p - data);
    *s_len = len;
    *s_data = p;
    return 0;
}

static esp_err_t esp_delta_ota_read_cb(uint8_t *buf_p, size_t size, int src_offset)
{
    if (size <= 0)
    {
        return ESP_ERR_INVALID_ARG;
    }
    return esp_partition_read(esp_ota_get_running_partition(), src_offset, buf_p, size);
}

static esp_err_t esp_delta_ota_write_cb(const uint8_t *buf_p, size_t size, void *user_data)
{
    if (size <= 0 || !user_data)
        return ESP_ERR_INVALID_ARG;
    return esp_ota_write((esp_ota_handle_t)user_data, buf_p, size);
}

uint8_t ota_app_bin_stream_process_cb(void *handle, uint8_t *data, uint32_t len)
{
    if (!handle)
        handle = ota_default_handle;
    ota_t ota = handle;
    if (!ota)
    {
        dbus_change_to_process("ota", "Missing parameter");
        return 2;
    }
    esp_err_t err;
    if (len && data)
    {
        if (!ota->status)
        {
            //  upgrade......
            // if (!(ota->type & OTA_TYPE_FILE))
            // {
            //     if (http_check_header(ota, &data, &len))
            //         return 2;
            // }
            ota->delta_ota = 0;
            uint16_t version = ota->last_version;
            uint32_t recv_magic = *(uint32_t *)data;
            if (ota->offset == 0)
            {
                const esp_partition_t *run_app_partition = esp_ota_get_running_partition();
                /* 差分OTA */
                if (recv_magic == 0xfccdde10 /* Delta OTA magic header */)
                {
                    ESP_LOGW(TAG, "Current OTA streams are differential upgrade packets");
                    uint8_t *digest = (uint8_t *)(data + 4);
                    uint8_t sha_256[32] = {0};
                    esp_partition_get_sha256(run_app_partition, sha_256);
                    if (memcmp(sha_256, digest, 32) != 0)
                    {
                        sprintf(ota->err_str, "Invalid patch for current firmware");
                        ESP_LOGE(TAG, ota->err_str);
                        dbus_change_to_process("ota", "Invalid patch for current firmware");
                        return 2;
                    }
                    /* Image Magic:4Byte + SHA256:32Byte + Reserved Header:92Byte = 128Byte */
                    data += 128;
                    len -= 128;

                    ota->write_size += 128;
                }
                else /* 完整包OTA */
                {
                    if (ota->buffer_offset + len < sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t) + sizeof(esp_app_desc_t))
                    {
                        if (ota->buffer)
                        {
                            len = len > ota->buffer_size ? ota->buffer_size : len;
                            memcpy(ota->buffer, data, len);
                            ota->buffer_offset = len;
                        }
                        return 0;
                    }

                    if (ota->buffer_offset)
                    {
                        memcpy(ota->buffer + ota->buffer_offset, data, len);
                        data = ota->buffer;
                        len = ota->buffer_offset + len;
                        ota->buffer_offset = 0;
                    }

                    const int app_desc_offset = sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t);
                    esp_app_desc_t *app_info = (esp_app_desc_t *)(data + app_desc_offset);
                    if (app_info->magic_word != ESP_APP_DESC_MAGIC_WORD)
                    {
                        ESP_LOG_BUFFER_HEXDUMP(TAG, data, len, ESP_LOG_INFO);
                        sprintf(ota->err_str, "The updated image is not an ESP-IDF image");
                        ESP_LOGE(TAG, "%s", ota->err_str);
                        dbus_change_to_process("ota", "The updated image is not an ESP-IDF image");
                        return 2;
                    }
                    esp_app_desc_t run_app_info = {0};
                    esp_ota_get_partition_description(run_app_partition, &run_app_info);
                    if (strncmp(app_info->project_name, run_app_info.project_name, sizeof(run_app_info.project_name)))
                    {
                        ESP_LOG_BUFFER_HEXDUMP(TAG, data, len, ESP_LOG_INFO);
                        sprintf(ota->err_str, "The updated firmware project name is inconsistent with the current firmware");
                        ESP_LOGE(TAG, "%s", ota->err_str);
                        dbus_change_to_process("ota", "The updated firmware project name is inconsistent with the current firmware");
                        ESP_LOGE(TAG, "Project name %s != %s", app_info->project_name, run_app_info.project_name);
                        return 2;
                    }

                    uint16_t len = strlen(strchr(app_info->version, '.') + 1);
                    if (len == 1)
                        len = 10;
                    else if (len == 2)
                        len = 100;
                    else if (len == 3)
                        len = 1000;
                    uint16_t version = atof(app_info->version) * len;
                    version = version / len << 8 | version % len;

                    uint8_t version_err = 0;
                    if (!ota->checkVersion)
                    {
                        if (version <= ota->version)
                            version_err = 1;
                    }
                    else if (ota->checkVersion == 1)
                    {
                        if (version < ota->version)
                            version_err = 1;
                    }
                    if (version_err)
                    {
                        sprintf(ota->err_str,
                                "The upgrade version is incorrect(%d.%d --> %d.%d)",
                                ((uint8_t *)&ota->version)[1], ((uint8_t *)&ota->version)[0],
                                ((uint8_t *)&version)[1], ((uint8_t *)&version)[0]);
                        ESP_LOGE(TAG, "%s", ota->err_str);
                        dbus_change_to_process("ota", "The latest version");
                        return 2;
                    }
                    ota->last_version = version;
                }
            }

            /* start */
            const esp_partition_t *update_partition = esp_ota_get_next_update_partition(NULL);
            if (!update_partition)
            {
                sprintf(ota->err_str, "There is only one OTA partition.");
                ESP_LOGW(TAG, "%s", ota->err_str);
                ota->status = OTA_STATUS_FAILED_FIND_PARTITION;
                dbus_change_to_process("ota", "There is only one OTA partition");
                return 2;
            }
            ESP_LOGI(TAG, "OTA update partition : %s", update_partition->label);
            ESP_LOGI(TAG, "length:%lu offset:%lu", ota->length, ota->offset);
            if (!ota->length)
                ota->length = update_partition->size;
            err = esp_ota_begin(update_partition, OTA_WITH_SEQUENTIAL_WRITES, &ota->update_handle);
            if (err != ESP_OK)
            {
                sprintf(ota->err_str, "esp_ota_begin failed (%s)", esp_err_to_name(err));
                ESP_LOGW(TAG, "%s", ota->err_str);
                ota->status = OTA_STATUS_FAILED_PARTITION_WRITE_BEGIN;
                dbus_change_to_process("ota", "Failed to erase partition");
                return 2;
            }
            if (recv_magic == 0xfccdde10 /* Delta OTA magic header */)
            {
                esp_delta_ota_cfg_t cfg = {
                    .read_cb = (src_read_cb_t)&esp_delta_ota_read_cb,
                    .user_data = (void *)ota->update_handle,
                    .write_cb_with_user_data = &esp_delta_ota_write_cb,
                };

                ota->delta_handle = esp_delta_ota_init(&cfg);
                if (ota->delta_handle == NULL)
                {
                    sprintf(ota->err_str, "esp_delta_ota_init failed");
                    ESP_LOGW(TAG, "%s", ota->err_str);
                    ota->status = OTA_STATUS_FAILED_PARTITION_WRITE;
                    dbus_change_to_process("ota", "esp_delta_ota_init failed");
                    return 2;
                }

                ota->delta_ota = 1;
            }
            else if (ota->offset) /* 偏移只支持完整更新 */
            {
                typedef struct ota_ops_entry_
                {
                    uint32_t handle;
                    const esp_partition_t *part;
                    bool need_erase;
                    uint32_t wrote_size;
                    uint8_t partial_bytes;
                    WORD_ALIGNED_ATTR uint8_t partial_data[16];
                    LIST_ENTRY(ota_ops_entry_)
                    entries;
                } ota_ops_entry_t;
                extern ota_ops_entry_t *get_ota_ops_entry(esp_ota_handle_t handle);
                ota_ops_entry_t *it = get_ota_ops_entry(ota->update_handle);
                it->wrote_size += ota->offset;
            }

            dbus_change_to_process("ota_name", "App");
            dbus_only_change("state.ota.run", 1);
            ota->status = OTA_STATUS_RUNNING;
        }

        if (ota->status == OTA_STATUS_RUNNING)
        {
            uint8_t failed = 0;
            if (ota->delta_ota)
            {
                if (esp_delta_ota_feed_patch(ota->delta_handle, (const uint8_t *)data, len) < 0)
                    failed = 1;
            }
            else if (esp_ota_write(ota->update_handle, (const void *)data, len) != ESP_OK)
                failed = 1;
            if (failed)
            {
                esp_ota_abort(ota->update_handle);
                sprintf(ota->err_str, "Failed to write ota ota->buffer to update partition");
                ESP_LOGW(TAG, "%s", ota->err_str);
                ota->status = OTA_STATUS_FAILED_PARTITION_WRITE;
                dbus_change_to_process("ota", "Failed to write partition");
                dbus_only_change("state.ota.run", 0);
                return 2;
            }
            ota->offset += len;
            ota->write_size += len;
            ota->progress = ota->write_size * 100 / ota->length;
            dbus_change_to_process("ota", ota->progress);
            LOG_PROGRESS("App write %ld %06ld %d%%", len, ota->write_size, ota->progress);
        }
    }
    else if (!len && !data)
    {
        if (ota->status == OTA_STATUS_NULL)
            return 0;
        /* end */
        if (ota->status == OTA_STATUS_FAILED_CALL_ABORT)
        {
            ESP_LOGW(TAG, "User call abort ota");
            esp_ota_abort(ota->update_handle);
            return 1;
        }
        ESP_LOGI(TAG, "close ota process.");
        // if (ota->write_size < ota->length)
        // {
        //     sprintf(ota->err_str, "The update size is incorrect");
        //     ESP_LOGW(TAG, "%s", ota->err_str);
        //     esp_ota_abort(ota->update_handle);
        //     dbus_change_to_process("ota", "The update size is incorrect");
        //     ota->update_handle = 0;
        //     ota->status = OTA_STATUS_FAILED_WRITE_ABORT;
        //     re = 2;
        //     goto ota_app_end;
        // }

        if (ota->delta_handle)
        {
            err = esp_delta_ota_finalize(ota->delta_handle);
            if (err != ESP_OK)
            {
                sprintf(ota->err_str, "esp_delta_ota_finalize failed (%s)!", esp_err_to_name(err));
                ESP_LOGE(TAG, "%s", ota->err_str);
                ota->status = OTA_STATUS_FAILED_ERR;
            }
            esp_delta_ota_deinit(ota->delta_handle);
            ota->delta_handle = NULL;
            ota->delta_ota = 0;
        }

        //  OTA upgrade end
        err = esp_ota_end(ota->update_handle);
        ota->update_handle = 0;
        ota->status = OTA_STATUS_NULL;
        // hal_global->hal.ota.states = OTA_STATUS_NULL;
        uint8_t re = 0;
        if (err != ESP_OK)
        {
            if (err == ESP_ERR_OTA_VALIDATE_FAILED)
            {
                sprintf(ota->err_str, "Image validation failed, image is corrupted");
                ota->status = OTA_STATUS_FAILED_VALIDATE;
                ESP_LOGE(TAG, "%s", ota->err_str);
            }
            else
            {
                sprintf(ota->err_str, "esp_ota_end failed (%s)!", esp_err_to_name(err));
                ota->status = OTA_STATUS_FAILED_ERR;
                ESP_LOGE(TAG, "%s", ota->err_str);
            }
            dbus_change_to_process("ota", "Failed to save partition");
            re = 2;
            goto ota_app_end;
        }
        else
        {
            const esp_partition_t *update_partition = esp_ota_get_next_update_partition(NULL);
            err = esp_ota_set_boot_partition(update_partition);
            if (err != ESP_OK)
            {
                sprintf(ota->err_str, "esp_ota_set_boot_partition failed (%s)!", esp_err_to_name(err));
                ota->status = OTA_STATUS_FAILED_SET_BOOT_APP;
                ESP_LOGE(TAG, "%s", ota->err_str);
                dbus_change_to_process("ota", "Failed to Boot partition");
                re = 1;
                goto ota_app_end;
            }
            ota->tick = esp_log_timestamp() + 5000;
        }
    ota_app_end:
        dbus_only_change("state.ota.offset", ota->offset);
        dbus_only_change("state.ota.run", 0);
        return re;
    }
    return 0;
}

uint8_t ota_www_bin_stream_process_cb(void *handle, uint8_t *data, uint32_t len)
{
    if (!handle)
        handle = ota_default_handle;
    ota_t ota = handle;
    if (!ota)
    {
        dbus_change_to_process("ota", "Missing parameter");
        return 2;
    }
    if (data && len)
    {
        if (!ota->status)
        {
            ota->par_info = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_FAT, "www");
            if (!ota->par_info)
                ota->par_info = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_LITTLEFS, "www");
            if (!ota->par_info)
            {
                ota->status = OTA_STATUS_FAILED_FIND_PARTITION;
                ESP_LOGE(TAG, "No partition was found(WWW)");
                dbus_change_to_process("ota", "Not find partition");
                return 2;
            }

            if (!(ota->type & OTA_TYPE_FILE))
            {
                // if (http_check_header(ota, &data, &len))
                //     return 2;
                if (ota->length >= 1024 * 1024)
                    ESP_LOGI(TAG, "www size : %2lu.%luMB", ota->length / 1024 / 1024, (ota->length / 1024 % 1024) * 10 / 1024);
                else if (ota->length >= 1024)
                    ESP_LOGI(TAG, "www size : %2lu.%luKB", ota->length / 1024, (ota->length % 1024) * 10 / 1024);
                else
                    ESP_LOGI(TAG, "www size : %luByte", ota->length);
                if (!ota->length)
                    ota->length = ota->par_info->size;
                // if (ota->length < 100 * 1024)
                // {
                //     ota->status = OTA_STATUS_FAILED_ERR;
                //     ESP_LOGE(TAG, "ota www length too small");
                //     dbus_change_to_process("ota", "WWW length too small");
                //     return 2;
                // }
            }
            ota->status = OTA_STATUS_RUNNING;
            // hal_global->hal.ota.states = OTA_STATUS_RUNNING;
            dbus_change_to_process("ota_name", "www");
            dbus_only_change("state.ota.run", 1);
        }

        // must erase the partition before writing to it
        uint32_t first_sector = ota->write_size / SPI_FLASH_SEC_SIZE;            // first affected sector
        uint32_t last_sector = (ota->write_size + len - 1) / SPI_FLASH_SEC_SIZE; // last affected sector
        esp_err_t ret = ESP_OK;
        if ((ota->write_size % SPI_FLASH_SEC_SIZE) == 0)
            ret = esp_partition_erase_range(ota->par_info, ota->write_size, ((last_sector - first_sector) + 1) * SPI_FLASH_SEC_SIZE);
        else if (first_sector != last_sector)
            ret = esp_partition_erase_range(ota->par_info, (first_sector + 1) * SPI_FLASH_SEC_SIZE, (last_sector - first_sector) * SPI_FLASH_SEC_SIZE);
        if (ret != ESP_OK)
            return ret;

        if (!esp_partition_write(ota->par_info, ota->write_size, data, len))
        {
            ota->write_size += len;
            ota->progress = ota->write_size * 100 / ota->length;
            // hal_global->hal.ota.states = OTA_STATUS_RUNNING;
            dbus_change_to_process("ota", ota->progress);
            LOG_PROGRESS("www write %ld %06ld %d%%", len, ota->write_size, ota->progress);
        }
        else
        {
            ota->status = OTA_STATUS_FAILED_PARTITION_WRITE;
            // hal_global->hal.ota.states = OTA_STATUS_NULL;
            dbus_change_to_process("ota", "Failed to write partition");
            dbus_only_change("state.ota.run", 0);
            return 2;
        }
    }
    else
    {
        if (ota->status == OTA_STATUS_FAILED_CALL_ABORT)
        {
            ESP_LOGW(TAG, "User call abort ota");
            return 1;
        }
        dbus_only_change("state.ota.run", 0);
        ota->status = OTA_STATUS_NULL;
        // hal_global->hal.ota.states = OTA_STATUS_NULL;
        ESP_LOGI(TAG, "The www upgrade end.");
    }
    return 0;
}
